# -*- coding: UTF-8 -*-
from captcha.image import ImageCaptcha  # pip install captcha
from PIL import Image
import random
import time
import settings
import os


def generate_captcha_text():
    captcha_text = []
    for i in range(settings.MAX_CAPTCHA):
        c = random.choice(settings.ALL_CHAR_SET)
        captcha_text.append(c)
    return ''.join(captcha_text)


def generate_captcha_text_and_image():
    image = ImageCaptcha()
    captcha_text = generate_captcha_text()
    captcha_image = Image.open(image.generate(captcha_text))
    return captcha_text, captcha_image


if __name__ == '__main__':
    # 训练集的大小
    train_count = 3000
    # 测试集的大小
    val_count = 500
    train_path = settings.TRAIN_DATASET_PATH
    val_path = settings.VAL_DATASET_PATH
    if not os.path.exists(val_path):
        os.makedirs(val_path)
    if not os.path.exists(train_path):
        os.makedirs(train_path)
    for i in range(train_count):
        now = str(int(time.time()))
        text, image = generate_captcha_text_and_image()
        filename = text + '_' + now + '.png'
        image.save(train_path + os.path.sep + filename)
        print('train_img is saved %d : %s' % (i + 1, filename))
    for i in range(val_count):
        now = str(int(time.time()))
        text, image = generate_captcha_text_and_image()
        filename = text + '_' + now + '.png'
        image.save(val_path + os.path.sep + filename)
        print('val_img is saved %d : %s' % (i + 1, filename))
